#include "tcp_sender.hh"

#include "tcp_config.hh"

#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&... /* unused */) {}

using namespace std;

#define DEBUG_SND 1

#if DEBUG_SND
#define PRINTF_SND(...) do { \
    printf(__VA_ARGS__); \
} while(0)
#else
#define PRINTF_SND(...)
#endif

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _stream(capacity) {
       rto = _initial_retransmission_timeout;
    }

uint64_t TCPSender::bytes_in_flight() const { 
    return (_next_seqno  - _ack_seqno); 
}

void TCPSender::fill_window() {
    
start:
    TCPSegment seg;
    
    bool allow_data_send = false;
    // min_sz > stream_in.buffer_size, the min will be select
    size_t min_sz = std::min(_win_size, static_cast<size_t>(TCPConfig::MAX_PAYLOAD_SIZE));
    //size_t max_sz = std::max(_win_size, static_cast<size_t>(TCPConfig::MAX_PAYLOAD_SIZE));
    //size_t min_sz = min(TCPConfig::MAX_PAYLOAD_SIZE, min(stream_in().buffer_size(), _win_size));
    
    PRINTF_SND("\n#fill_window: ack %lu, win_sz %lu,  min_sz %lu, next_seq %lu eof %d\n", \
                                _ack_seqno, _win_size, min_sz, _next_seqno, stream_in().eof());
    // "ack + windows_size" is the right edge for _next_seqno
    if ((min_sz + _next_seqno) > (_ack_seqno + _win_size)) {
        min_sz = (_ack_seqno + _win_size) <  _next_seqno ? 0 : (_ack_seqno + _win_size -  _next_seqno);       
    }
    
#if 1
    // only for window size == 0
    // after SYN_ACKED, when _window_sz is 0, set as 1
    if ((_win_size == 0) && (next_seqno_absolute() > bytes_in_flight()) && !stream_in().eof()  &&(min_sz == 0)) {
        min_sz = 1;
        PRINTF_SND("*SYN_ACKED min_sz %lu\n", min_sz);
        allow_data_send = true;
    }
#endif
    
    if (_win_size && (min_sz == 0)) {
        PRINTF_SND("windows is full: min_sz %lu\n", min_sz);
        return;
    }

    // CLOSED
    if(next_seqno_absolute() == 0) {
        // how to choose _isn ??? fixed 
        // send SYN & enquque stream out
        //_next_seqno = 1;
        
        //header.syn = true;
        seg.header().syn = true;
        
        PRINTF_SND("CLOSED: set SYN\n");
        goto send;
    }


    // SYN_SENT
    // SYN sentbut not ack yet, don't send data but can send fin 
    if ((next_seqno_absolute() > 0) && (next_seqno_absolute() == bytes_in_flight())) {
        PRINTF_SND("SYN_SENT: but not ack\n");
        return;
    }

    // SYN_ACKED, can send data
    if ((next_seqno_absolute() > bytes_in_flight()) && !stream_in().eof()) {
        allow_data_send = true;
        PRINTF_SND("SYN_ACKED: can data next_seqno %lu, bytes_fly %lu, eof %d\n", \
                                     next_seqno_absolute(), bytes_in_flight(), stream_in().eof());
        
        goto send;
    }

        
    // SYN_ACKED, can send FIN
    if (stream_in().eof() && 
        (next_seqno_absolute() < stream_in().bytes_written() + 2)) {
        
        // if remain data > (win_sz & payload_len), send data first
        // set FIN after remain data is short enougph
        size_t remain_data_len = _stream.bytes_written() - _stream.bytes_read();
        if (remain_data_len <= min_sz) {
            seg.header().fin = true;
        }
        
        allow_data_send = true;
        
        // only for window_size == 0
        if (!_win_size && !min_sz) {
            min_sz = 1;
        }

        PRINTF_SND("SYN_ACKED: set FIN, remain %lu, min_sz %lu, FIN %d\n", remain_data_len, min_sz, seg.header().fin);
        goto send;
    }

    if (stream_in().eof() && 
        (next_seqno_absolute() == stream_in().bytes_written() + 2)) {
        // FIN_ACKED, connection closed totally
        if (bytes_in_flight() == 0) {
            return ;   
        } else { // FIN_SENT | sender.bytes_in_flight() > 0, some ACK not recv yet
            
            return ;   
        }
    }
    // fetch data from stream_in min(PAYLOAD,window_size)
    // generate tcp_seg & enqueque to stream_out queue ?
    // increase bytes_in_flight ?

send:

    std::string data;
    if (allow_data_send) {
        PRINTF_SND("do send: min_sz %lu.\n", min_sz);
        // fetch & erase data
        data =  _stream.peek_output(min_sz);
        size_t data_sz = data.length();
        if (data_sz == 0 && (!seg.header().syn && !seg.header().fin)) {
        
            PRINTF_SND("data len eq 0.\n");
            return;
        }
  
        PRINTF_SND("do send: ack %lu, win_sz %lu,  min_sz %lu, next_seq %lu eof %d\n", \
                                _ack_seqno, _win_size, min_sz, _next_seqno, stream_in().eof());
 #if 1
        if ((_win_size == 0) && \
            ((_ack_seqno + min_sz) <= _next_seqno)) {
            PRINTF_SND("!no remote win space, stop send!\n");
            return;
        }
#endif
        // TODO: pop when ack ???
        _stream.pop_output(data_sz);
        
        // stream  eof here and need win_size is enough
        if (stream_in().eof() && ((_ack_seqno + _win_size) >= (stream_in().bytes_written() + 2))) {
            PRINTF_SND("FIN: win_size enough & all writen.\n");
            seg.header().fin = true;
        } 

        // push to segment
        Buffer buffer{std::string(data)};
        seg.payload() = buffer;
    }  

    // fill_window  don't do retry 
    seg.header().seqno = next_seqno();
    
    // push to sent list 
    SentSeg snd_seg(seg, 0, _initial_retransmission_timeout, _next_seqno, _ticks);
    _sent_list.push_back(snd_seg);

    // update next seqno
    _next_seqno += seg.length_in_sequence_space();
   
    PRINTF_SND("LABEL send: _next_seqno %lu, seg_len %lu, allow_data_send %d, data len %ld, eof %d\n", \
                            _next_seqno,seg.length_in_sequence_space(), allow_data_send, data.length(), stream_in().eof());

    _segments_out.push(seg);

    
    if (!stream_in().buffer_empty() && (_next_seqno < (_ack_seqno + _win_size))) {
        PRINTF_SND("@re-fill: stream_in sz %lu, next_seq %lu, ack + win %lu\n", \
                              stream_in().buffer_size(), _next_seqno, (_ack_seqno + _win_size) );
        goto start;
    }

}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) { 
    uint64_t ack= unwrap(ackno, _isn, _next_seqno);
    if (ack > _next_seqno) {
        PRINTF_SND("ack_recv: invalid ack %lu, exp _next_seqno %lu\n", ack, _next_seqno);
        return;
    }

    if (ack > _ack_seqno) {
        _ack_seqno = ack;   
        
        // reset 
        cons_retx_cnt = 0;
        rto = _initial_retransmission_timeout;
        
        // sync all seq last seqno when new data recved
        for (auto& seg : _sent_list) {
            seg.last_snd_time = _ticks;
        }
    }   
    _win_size = window_size;
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) { 
    _ticks += ms_since_last_tick;
    
    bool retx = false;

#if 0
    // clean seqno <= _ack_seqno in vector
    _sent_list.erase(std::remove_if(_sent_list.begin(), _sent_list.end(),
                         [this](SentSeg seg) { 
                                if (0 == _win_size) {
                                    PRINTF_SND("win_sz 0: _ack %lu seq %lu\n", _ack_seqno, seg.seqno);
                                    return false;
                                }
                                PRINTF_SND("CLEAN: _ack %lu seq %lu, seq_space_len %lu\n", \
                                                   _ack_seqno, seg.seqno, seg.seg.length_in_sequence_space());
                                
                                //if (seg.seg.header().fin && )
                                return seg.seqno < _ack_seqno; 
                                //return (seg.seqno + seg.seg.length_in_sequence_space()) < _ack_seqno; 
                            }),
          _sent_list.end());
#endif

    // RETX check
    for (auto& seg : _sent_list) {
        
        PRINTF_SND("\n>retx segno %lu\n", seg.seqno);
        if ((_ticks - seg.last_snd_time) < rto) {
            PRINTF_SND("PASS1: _ticks %lu, since last %lu, seg last snd %lu, RTO %u\n",_ticks, ms_since_last_tick, seg.last_snd_time, rto);
            continue;
        }
#if 1
        //if (_ack_seqno > (seg.seqno + seg.seg.length_in_sequence_space()) ) {
        if (_ack_seqno != seg.seqno ) {
            // the last seg is special @?
            if (seg.seg.header().fin) {
                PRINTF_SND("LAST1: ack %lu > (seg seqno %lu, seg space len %lu) RTO %u\n", \
                               _ack_seqno, seg.seqno, seg.seg.length_in_sequence_space(),  rto);
            } else {
                PRINTF_SND("PASS2: ack %lu > (seg seqno %lu, seg space len %lu) RTO %u\n", \
                               _ack_seqno, seg.seqno, seg.seg.length_in_sequence_space(),  rto);
                continue;
            }
        }
#endif  
        PRINTF_SND("RETX: _ticks %lu,sin last %lu, last snd %lu, RTO %u, _ack_seqno %lu, segno %lu\n", \
                          _ticks, ms_since_last_tick, seg.last_snd_time, rto, _ack_seqno, seg.seqno);
       
        seg.last_snd_time = _ticks;
       
       #if 0
        if (seg.retx_cnt < TCPConfig::MAX_RETX_ATTEMPTS) {
            seg.rto = seg.rto * 2;
            seg.retx_cnt++;
        }
        #endif
       _segments_out.push(seg.seg);
       
       retx = true;
        
       // ??? only send the oldest one ?? first match !
       break;
    } // loop end
    
    // really has retx operation && less than max attempt
    if ( retx && _win_size) { //&& (cons_retx_cnt < TCPConfig::MAX_RETX_ATTEMPTS)) {
        if (cons_retx_cnt < TCPConfig::MAX_RETX_ATTEMPTS) {
            rto = rto * 2; 
        }

        // statistics no limit, but max rto is limit by MAX_RETX_ATTEMPTS
        cons_retx_cnt++;
    }

#if 1
    // clean seqno <= _ack_seqno in vector
    _sent_list.erase(std::remove_if(_sent_list.begin(), _sent_list.end(),
                         [this](SentSeg seg) { 
                                if (0 == _win_size) {
                                    PRINTF_SND("win_sz 0: _ack %lu seq %lu\n", _ack_seqno, seg.seqno);
                                    return false;
                                }
                                PRINTF_SND("CLEAN: _ack %lu seq %lu, seq_space_len %lu\n", \
                                                   _ack_seqno, seg.seqno, seg.seg.length_in_sequence_space());
                                return seg.seqno < _ack_seqno; 
                                //return (seg.seqno + seg.seg.length_in_sequence_space()) < _ack_seqno; 
                            }),
          _sent_list.end());
#endif

}

unsigned int TCPSender::consecutive_retransmissions() const { return cons_retx_cnt; }

void TCPSender::send_empty_segment() {}
